#include <map>
#include <string>
#include <vector>
#include <memory>
#include <cstdio>
#include <functional>

class TreeNode;
using TreeNodePtr = std::shared_ptr<TreeNode>;

using TreeNodeVisitFunc = std::function<void(const std::string &k, const TreeNodePtr &v)>;

class TreeNode {
public:
    TreeNodePtr GetNode(std::string s)
    {
        if (m == nullptr) {
            m = std::make_unique<std::map<std::string, TreeNodePtr>>();
        }
        if (!HasNode(s)) {
            (*m)[s] = std::make_shared<TreeNode>();
        }
        return (*m)[s];
    }

    bool HasNode(std::string s)
    {
        return m->find(s) != m->end();
    }

    void Foreach(std::string pre, TreeNodeVisitFunc func) const
    {
        if (m == nullptr) {
            return;
        }
        for (const auto &[k, v] : *m) {
            func(pre + "." + k, v);
        }
    }

    int v = 0;

private:
    std::unique_ptr<std::map<std::string, TreeNodePtr>> m = nullptr;
};

void PrintNodes(const std::string &k, const TreeNodePtr &v)
{
    printf("%s %d\n", k.c_str(), v->v);
    v->Foreach(k, PrintNodes);
}

void A(const TreeNodePtr &node, std::function<void(const TreeNodePtr&)> func)
{
    TreeNodeVisitFunc getV = [&](const std::string &k, const TreeNodePtr &v) {
        func(v);
        v->Foreach(k, getV);
    };
    node->Foreach("root", getV);
}

void B(const TreeNodePtr &root)
{
    A(root, [&](const TreeNodePtr &node) {
        printf("B: %d\n", node->v);
    });
}

int main()
{
    const TreeNodePtr root = std::make_shared<TreeNode>();

    root->GetNode("a")->GetNode("b")->v = 1;
    root->GetNode("a")->v = 2;
    root->GetNode("b")->v = 3;
    root->GetNode("c")->v = 4;

    TreeNodePtr a = root;
    printf("%d\n", a->v);

    std::vector<std::string> query({"a", "b", "c"});
    for (const auto &q : query) {
        a = a->GetNode(q);
        printf("%d\n", a->v);
    }

    PrintNodes("root", root);

    std::vector<int> ints;
    TreeNodeVisitFunc getV = [&](const std::string &k, const TreeNodePtr &v) {
        ints.push_back(v->v);
        v->Foreach(k, getV);
    };
    root->Foreach("root", getV);

    for (const auto &i : ints) {
        printf("i: %d\n", i);
    }

    B(root);
    return 0;
}
